{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/fireworks-ai/cookbook/blob/main/integrations/Klavis/Use_Klavis_with_Fireworks.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "# Fireworks AI + Klavis AI Integration\n",
    "\n",
    "# <img src=\"../../static/fireworks-klavis.png\" width=\"500\">\n",
    "\n",
    "\n",
    "In this tutorial, we'll explore how to build an AI agent that integrates Fireworks AI's LLM with Klavis MCP Servers:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "## Prerequisites\n",
    "\n",
    "Before we begin, you'll need:\n",
    "\n",
    "- **Fireworks AI API key** - see here [fireworks.ai](https://fireworks.ai/)\n",
    "- **Klavis API key** - see here [klavis.ai](https://klavis.ai/)\n",
    "\n",
    "Make sure to keep these API keys secure and never commit them to version control!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.0\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.1.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "# Install the required packages\n",
    "%pip install -qU fireworks-ai klavis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from fireworks.client import Fireworks\n",
    "from klavis import Klavis\n",
    "from klavis.types import McpServerName, ToolFormat\n",
    "\n",
    "# Set environment variables\n",
    "os.environ[\"FIREWORKS_API_KEY\"] = \"fw_XXXXXXXXXXXXXX\"  # Replace with your actual Fireworks API key\n",
    "os.environ[\"KLAVIS_API_KEY\"] = \"XXXXXXXXXXXXX\"    # Replace with your actual Klavis API key\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "## Create AI Agent with MCP Integration\n",
    "\n",
    "Now we'll create an intelligent agent that can use MCP servers through Klavis API. This agent will:\n",
    "\n",
    "1. **Discover Tools**: Automatically find available tools from MCP servers\n",
    "2. **Function Calling**: Use Fireworks AI's function calling capabilities\n",
    "3. **Tool Execution**: Execute tools through Klavis API\n",
    "4. **Smart Responses**: Generate intelligent responses based on tool results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent:\n",
    "    def __init__(self, fireworks_client, klavis_client, mcp_server_url):\n",
    "        self.fireworks = fireworks_client\n",
    "        self.klavis = klavis_client\n",
    "        self.mcp_server_url = mcp_server_url\n",
    "        self.model = \"accounts/fireworks/models/qwen2p5-72b-instruct\"\n",
    "        print(f\"\ud83e\udd16 Agent initialized with model: {self.model}\")\n",
    "    \n",
    "    def process_request(self, user_message):\n",
    "        # 1. Get available tools\n",
    "        mcp_tools = self.klavis.mcp_server.list_tools(\n",
    "            server_url=self.mcp_server_url,\n",
    "            format=ToolFormat.OPENAI,\n",
    "        )\n",
    "        \n",
    "        # 2. Call LLM with tools\n",
    "        messages = [\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "            {\"role\": \"user\", \"content\": user_message}\n",
    "        ]\n",
    "        \n",
    "        response = self.fireworks.chat.completions.create(\n",
    "            model=self.model,\n",
    "            messages=messages,\n",
    "            tools=mcp_tools.tools\n",
    "        )\n",
    "        \n",
    "        assistant_message = response.choices[0].message\n",
    "        messages.append(assistant_message)\n",
    "        \n",
    "        # 3. If LLM wants to use tools\n",
    "        if assistant_message.tool_calls:\n",
    "            \n",
    "            # Execute each tool call\n",
    "            for tool_call in assistant_message.tool_calls:\n",
    "                tool_name = tool_call.function.name\n",
    "                tool_args = json.loads(tool_call.function.arguments)\n",
    "                \n",
    "                print(f\"\ud83d\udee0\ufe0f Calling tool: {tool_name} with args: {tool_args}\")\n",
    "                # Call tool via Klavis SDK\n",
    "                tool_result = self.klavis.mcp_server.call_tools(\n",
    "                    server_url=self.mcp_server_url,\n",
    "                    tool_name=tool_name,\n",
    "                    tool_args=tool_args,\n",
    "                )\n",
    "                \n",
    "                messages.append({\n",
    "                    \"role\": \"tool\",\n",
    "                    \"tool_call_id\": tool_call.id,\n",
    "                    \"content\": str(tool_result)\n",
    "                })\n",
    "            \n",
    "            # 4. Get final response from LLM\n",
    "            final_response = self.fireworks.chat.completions.create(\n",
    "                model=self.model,\n",
    "                messages=messages\n",
    "            )\n",
    "            return final_response.choices[0].message.content\n",
    "        \n",
    "        # If no tools needed, return the assistant message directly\n",
    "        return assistant_message.content\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use Case 1: Summarize YouTube Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u2705 Created YouTube MCP instance\n",
      "\ud83e\udd16 Agent initialized with model: accounts/fireworks/models/qwen2p5-72b-instruct\n",
      "\ud83d\udee0\ufe0f Calling tool: get_youtube_video_transcript with args: {'url': 'https://www.youtube.com/watch?v=kPXvf2-C_Hs'}\n",
      "### Summary of the Video: \"How This AI Startup Grew by 100x in Just 6 Months | Fireworks AI, Lin Qiao\"\n",
      "\n",
      "**Title:** How This AI Startup Grew by 100x in Just 6 Months | Fireworks AI, Lin Qiao  \n",
      "**Published:** August 20, 2024  \n",
      "**Channel:** EO  \n",
      "**Duration:** 11 minutes 55 seconds  \n",
      "**Views:** 77,278  \n",
      "**Likes:** 1,650  \n",
      "**Comments:** 51\n",
      "\n",
      "#### **0:00 - 1:00: Introduction**\n",
      "- Introduction to Fireworks AI, one of the fastest-growing AI startups in the industry.\n",
      "- Over the past six months, Fireworks AI has seen a 100x increase in traffic.\n",
      "- In July, the company closed a Series B funding round led by Sequoia, bringing its valuation to $552 million.\n",
      "\n",
      "#### **1:00 - 3:30: Lin Qiao's Background**\n",
      "- Lin Qiao, the founder of Fireworks AI, previously served as the Head of PyTorch.\n",
      "- She shares how her experiences at PyTorch inspired the creation of Fireworks AI.\n",
      "- Lin highlights the importance of staying connected with the latest AI developments and the community.\n",
      "\n",
      "#### **3:30 - 5:00: The AI Transformation**\n",
      "- Discussion on the ongoing AI transformation and its impact on various industries.\n",
      "- Lin emphasizes the need for adaptability and continuous learning in the AI landscape.\n",
      "- She talks about the rapid pace of innovation and the importance of staying ahead of the curve.\n",
      "\n",
      "#### **5:00 - 7:00: Fireworks AI's Growth Strategy**\n",
      "- Explanation of the key strategies that drove Fireworks AI's exponential growth.\n",
      "- Focus on user-centric design and scalable infrastructure.\n",
      "- Importance of community engagement and feedback loops in product development.\n",
      "\n",
      "#### **7:00 - 9:00: Insights for Software Engineers**\n",
      "- Lin provides insights into what it takes to excel as a software engineer in the AI industry.\n",
      "- Emphasis on practical skills, theoretical knowledge, and the ability to work in interdisciplinary teams.\n",
      "- Advice on building a strong foundation in AI and staying updated with the latest research.\n",
      "\n",
      "#### **9:00 - 11:00: The Future of Fireworks AI**\n",
      "- Lin shares her vision for the future of Fireworks AI.\n",
      "- Plans for expanding into new markets and developing cutting-edge AI applications.\n",
      "- The importance of maintaining a startup culture even as the company scales.\n",
      "\n",
      "**Contact Information:**\n",
      "- **Twitter:** @EO__Global\n",
      "- **LinkedIn:** @EO STUDIO\n",
      "- **Instagram:** @eostudio.official\n",
      "\n",
      "**Subtitles:**\n",
      "- Created using [XL8.ai](http://xl8.ai/) machine translation.\n",
      "\n",
      "This video provides valuable insights into the success of Fireworks AI and the broader AI industry, making it a must-watch for entrepreneurs, software engineers, and anyone interested in AI's future.\n"
     ]
    }
   ],
   "source": [
    "YOUTUBE_VIDEO_URL = \"https://www.youtube.com/watch?v=kPXvf2-C_Hs\" # pick a video you like!\n",
    "\n",
    "# 1. Initialize Fireworks AI client and Klavis client\n",
    "fireworks_client = Fireworks(api_key=os.getenv(\"FIREWORKS_API_KEY\"))\n",
    "klavis_client = Klavis(api_key=os.getenv(\"KLAVIS_API_KEY\"))\n",
    "\n",
    "# 2. Create YouTube MCP server instance\n",
    "youtube_mcp_instance = klavis_client.mcp_server.create_server_instance(\n",
    "    server_name=McpServerName.YOUTUBE,\n",
    "    user_id=\"1234\",\n",
    ")\n",
    "\n",
    "# 3. Create an agent with YouTube MCP server\n",
    "agent = Agent(fireworks_client, klavis_client, youtube_mcp_instance.server_url)\n",
    "\n",
    "response = agent.process_request(\n",
    "    f\"Summarize this YouTube video with timestamps: {YOUTUBE_VIDEO_URL}\"\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use Case 2: Send Email\n",
    "\n",
    "**Note**: Gmail integration requires OAuth authentication, so you'll need to authorize the application in your browser. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import webbrowser\n",
    "\n",
    "# Create Gmail MCP server instance\n",
    "gmail_mcp_instance = klavis_client.mcp_server.create_server_instance(\n",
    "    server_name=McpServerName.GMAIL,\n",
    "    user_id=\"1234\",\n",
    ")\n",
    "\n",
    "# Redirect to Gmail OAuth page\n",
    "webbrowser.open(gmail_mcp_instance.oauth_url)\n",
    "print(f\"\ud83d\udd10 Opening OAuth authorization for Gmail, if you are not redirected, please open the following URL in your browser: {gmail_mcp_instance.oauth_url}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The email has been sent successfully to zihaolin@klavis.ai with the subject \"Hello, World!\" and the body \"This is a test email.\" The email ID is 1977620de59daa96.\n"
     ]
    }
   ],
   "source": [
    "EMAIL_SUBJECT = \"Hello, World!\"\n",
    "EMAIL_BODY = \"This is a test email.\"\n",
    "EMAIL_RECIPIENT = \"zihaolin@klavis.ai\" # replace with your email\n",
    "\n",
    "# Create an agent with Gmail MCP server\n",
    "agent = Agent(fireworks_client, klavis_client, gmail_mcp_instance.server_url)\n",
    "\n",
    "response = agent.process_request(\n",
    "    f\"Send an email to {EMAIL_RECIPIENT} with subject {EMAIL_SUBJECT} and body {EMAIL_BODY}\"\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "## Summary\n",
    "\n",
    "This tutorial demonstrated how to integrate Fireworks AI's function calling capabilities with Klavis MCP servers to create powerful AI applications. We covered two practical examples:\n",
    "\n",
    "**\ud83c\udfa5 YouTube Integration**: Built an AI assistant that can automatically summarize YouTube videos by extracting transcripts and providing detailed, timestamped summaries.\n",
    "\n",
    "**\ud83d\udce7 Gmail Integration**: Created an AI-powered email assistant that can send emails through Gmail with OAuth authentication.\n",
    "\n",
    "### Key Takeaways:\n",
    "- **Easy Setup**: Klavis MCP servers can be created with just a few lines of code using the official SDK\n",
    "- **Fireworks AI Compatible**: All tools are formatted for seamless Fireworks AI function calling\n",
    "- **Versatile**: Support for both simple APIs (YouTube) and OAuth-authenticated services (Gmail)\n",
    "- **Scalable**: The same pattern can be applied to any of the 100+ MCP servers available in Klavis\n",
    "\n",
    "### \ud83d\ude80 Next Steps\n",
    "- **Explore More MCP Servers**: Try other available servers like Slack, Notion, CRM, etc.\n",
    "- **Try Different Fireworks AI Models**: Experiment with various models like Llama, Mixtral, or Deepseek for different use cases\n",
    "- **Build Complex Multi-Server Workflows**: Create sophisticated agents that combine Gmail + Slack + Notion for complete business automation\n",
    "- **Production Deployment**: Scale these patterns for production applications\n",
    "\n",
    "### \ud83d\udd17 Useful Resources\n",
    "- [Fireworks AI Documentation](https://docs.fireworks.ai/)\n",
    "- [Klavis AI Documentation](https://www.klavis.ai/docs/)\n",
    "- [MCP Protocol Specification](https://modelcontextprotocol.io/)\n",
    "\n",
    "**Happy building!** \ud83d\ude80\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}